import numpy as np
import image_io
from image_utils import logging
import os
import cv2

def split_image_channels(image):
    '''
    该函数用于分割图像的颜色通道
    :param image: 输入图像，numpy数组
    :return: 图像的颜色通道列表
    '''
    try:
        logging.info("开始分割图像颜色通道")
        image = cv2.imread('path_to_your_image.jpg')
        
        # TODO: 1. 检测多通道图像 image 是否真多通道的，如果是，显示错误信息并返回
        # 提示: 使用 raise 抛出异常信息：图像通道数必须大于等于2
        def check_multichannel_image(image):
            # 获取图像的通道数
            channels = image.shape[2] if len(image.shape) > 2 else 1

            # 检查通道数是否大于等于2
            if channels < 2:
                # 抛出异常，显示错误信息
                raise ValueError("图像通道数必须大于等于2。")

                # 如果通道数足够，则继续执行后续操作
            # 这里只是返回，因为题目没有要求其他操作
            return

            # 假设 image 是您已经读取的多通道图像

        # image = cv2.imread('path_to_your_image.jpg')

        # 调用函数检查图像
        try:
            check_multichannel_image(image)
            print("图像是多通道的，没有错误。")
        except ValueError as e:
            print(e)
        
        # TODO: 2. 使用numpy的split函数，沿着通道维度分割图像
        # 确保图像是多通道的
        if len(image.shape) == 3:
            # 获取通道数
            channels = image.shape[2]

            # 沿着通道维度分割图像
            # OpenCV中图像通道的顺序是BGR，因此分割得到的数组顺序也是BGR
            channel_images = np.split(image, channels, axis=2)

            # 现在channel_images是一个包含三个NumPy数组的列表，每个数组代表一个颜色通道
            blue_channel, green_channel, red_channel = channel_images

            # 输出每个通道的形状，确认分割成功
            print("Blue channel shape:", blue_channel.shape)
            print("Green channel shape:", green_channel.shape)
            print("Red channel shape:", red_channel.shape)
        else:
            print("图像不是多通道的，无法进行分割。")
        
        # TODO: 3. 把分割后的通道图像添加到一个列表 channel_list 中，注意要去掉多余的维度
        # 确保图像是多通道的
        if len(image.shape) == 3:
            # 获取通道数
            channels = image.shape[2]

            # 沿着通道维度分割图像
            channel_images = np.split(image, channels, axis=2)

            # 使用列表推导式和np.squeeze去掉每个通道图像的多余维度
            channel_list = [np.squeeze(channel) for channel in channel_images]

            # 现在channel_list包含没有多余维度的通道图像
            print("Channel list created with squeezed channel images.")
        else:
            print("图像不是多通道的，无法进行分割。")
        
        logging.info(f"成功分割图像颜色通道")
        return channel_list
    except Exception as e:
        # print(f"在分割图像颜色通道时发生错误: {e}")
        raise ValueError(f"在分割图像颜色通道时发生错误: {e}")

def merge_image_channels(file_paths):
    '''
    该函数用于合并多个单通道图像，创建一个多通道图像
    :param file_paths: 包含单通道图像文件路径的列表
    :return: 合并后的多通道图像
    '''
    try:
        logging.info("开始合并图像通道")
        
        # TODO: 想办法加载每个图像文件，然后将它们合并成一个多通道图像 merged_image
        # 注意: 检查每个图像是否是单通道的，如果不是，raise 一个异常信息: 要合并子图像必须是单通道图像，不能是多通道图像
        # 假设所有子图像都在同一个文件夹中，并且文件名遵循一定的模式
        folder_path = 'path_to_your_images_folder'
        image_files = [f for f in os.listdir(folder_path) if f.endswith('.png') or f.endswith('.jpg')]

        # 初始化一个空列表来存储加载的子图像
        channel_images = []

        # 加载每个子图像并检查其通道数
        for image_file in image_files:
            image_path = os.path.join(folder_path, image_file)
            image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)  # 确保以灰度模式加载

            # 检查图像是否是单通道的
            if image.ndim != 2:
                raise ValueError(f"要合并的子图像必须是单通道图像，不能是多通道图像。文件: {image_file}")

                # 将加载的单通道图像添加到列表中
            channel_images.append(image)

            # 检查是否加载了正确数量的通道图像
        if len(channel_images) < 3:
            raise ValueError("至少需要三个单通道图像来创建一个多通道图像。")

            # 确保所有子图像具有相同的尺寸
        height, width = channel_images[0].shape
        for image in channel_images:
            if image.shape != (height, width):
                raise ValueError("所有子图像必须具有相同的尺寸。")

                # 使用numpy的dstack函数沿着第三个维度堆叠图像以创建多通道图像
        merged_image = np.dstack(channel_images)

        # 输出合并后的图像的形状，以确认其是一个多通道图像
        print("Merged image shape:", merged_image.shape)
        
        logging.info("成功合并图像通道")
        return merged_image
    except Exception as e:
        raise ValueError(f"在合并图像通道时发生错误: {e}")